package org.example.com.zwh.util;

import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.util.Collections;

@Component // 添加Spring组件注解
public class RateLimiterUtil {
    private final RedisTemplate<String, Object> redisTemplate;
    // 令牌桶算法实现
    private static final String LUA_SCRIPT =
            "local key = KEYS[1]\n" +
                    "local max = tonumber(ARGV[1])\n" +
                    "local timeout = tonumber(ARGV[2])\n" +
                    "local now = redis.call('TIME')[1]\n" +
                    "local last_refill = tonumber(redis.call('HGET', key, 'last_refill') or 0)\n" +
                    "local tokens = tonumber(redis.call('HGET', key, 'tokens') or max)\n" +

                    // 计算需要补充的令牌
                    "local refill_tokens = math.floor((now - last_refill) / timeout)\n" +
                    "local new_tokens = math.min(max, tokens + refill_tokens)\n" +

                    // 更新令牌和时间
                    "if refill_tokens > 0 then\n" +
                    "   redis.call('HSET', key, 'last_refill', now)\n" +
                    "   redis.call('HSET', key, 'tokens', new_tokens)\n" +
                    "end\n" +

                    // 检查令牌是否足够
                    "if new_tokens >= 1 then\n" +
                    "   redis.call('HINCRBY', key, 'tokens', -1)\n" +
                    "   return 1\n" +
                    "else\n" +
                    "   return 0\n" +
                    "end";
    // 通过构造器注入RedisTemplate
    public RateLimiterUtil(RedisTemplate<String, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    // ... 保持原有LUA脚本不变 ...

    // 添加默认参数方法
    public boolean trySeckillLimit(String userId) {
        return tryAcquire("seckill:limit:" + userId, 3, 1);
    }

    // 保持原有tryAcquire方法不变，补充文档注释
    /**
     * 令牌桶限流核心方法
     * @param key      限流唯一键（如：userId）
     * @param max      最大令牌数（突发容量）
     * @param timeout  令牌补充间隔（秒）
     * @return 是否允许请求
     */
    public boolean tryAcquire(String key, int max, int timeout) {
        RedisScript<Long> script = new DefaultRedisScript<>(LUA_SCRIPT, Long.class);

        // timeout参数单位：秒
        Long result = redisTemplate.execute(
                script,
                Collections.singletonList(key),
                String.valueOf(max),
                String.valueOf(timeout)
        );

        return result != null && result == 1L;
    }
}